"""
本文件提供 根据混淆矩阵，计算出各大分类模型的评价指标的方法

"""
import os

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

"""
根据混淆矩阵计算出模型在总体的分类准确率
"""


def get_acc_total(confusion_matrix):
    acc_total = np.sum(np.diag(confusion_matrix)) / np.sum(confusion_matrix)
    return round(acc_total, 3)


"""
根据混淆矩阵计算出模型在每一个类的分类准确率
"""


def get_acc_per_class(confusion_matrix):
    acc_per_class = []
    for i in range(confusion_matrix.shape[0]):
        acc_per_class.append(confusion_matrix[i, i] / np.sum(confusion_matrix[:, i]))
    acc_per_class = [round(i, 3) for i in acc_per_class]
    return acc_per_class


"""
根据混淆矩阵计算出模型在每一个类的分类召回率
"""


def get_recall_per_class(confusion_matrix):
    recall_per_class = []
    for i in range(confusion_matrix.shape[0]):
        if np.sum(confusion_matrix[i, :] != 0):
            recall_per_class.append(confusion_matrix[i, i] / np.sum(confusion_matrix[i, :]))
        else:
            recall_per_class.append(0)
    recall_per_class = [round(i, 3) for i in recall_per_class]
    return recall_per_class


"""
根据混淆矩阵计算出模型在每一个类的分类F1值
"""


def get_f1_per_class(confusion_matrix):
    f1_per_class = []
    for i in range(confusion_matrix.shape[0]):
        f1_per_class.append(
            2 * confusion_matrix[i, i] / (np.sum(confusion_matrix[i, :]) + np.sum(confusion_matrix[:, i])))
    # f1_per_class保留小数点后3位
    f1_per_class = [round(i, 3) for i in f1_per_class]
    return f1_per_class


"""
获取输入目录下的 csv文件 混淆矩阵已经储存为csv文件
示例：
,Class_0,Class_1,Class_2,Class_3,Class_4
Class_0,9,0,0,0,0
Class_1,0,12,0,8,0
Class_2,0,2,10,1,1
Class_3,0,0,0,18,0
Class_4,0,0,1,2,6
"""


def get_confusion_all(input_csv):
    df = pd.read_csv(input_csv, index_col=0)
    class_names = df.columns[0:]
    confusion_matrix = np.zeros((len(class_names), len(class_names)))
    for i in range(len(class_names)):
        for j in range(len(class_names)):
            confusion_matrix[i, j] = df[class_names[i]][class_names[j]]
    acc_total = get_acc_total(confusion_matrix)
    acc_per_class = get_acc_per_class(confusion_matrix)
    recall_per_class = get_recall_per_class(confusion_matrix)
    f1_per_class = get_f1_per_class(confusion_matrix)
    return acc_total, acc_per_class, recall_per_class, f1_per_class


# a = get_confusion_matrix(r"日志数据/混淆矩阵/OurModel_confusion_matrix.csv")
# print(a)
# 获取混淆矩阵下所有的csv文件。
# 获取混淆矩阵下所有的csv文件。
for file_name in os.listdir("日志数据/混淆矩阵"):
    if file_name.endswith(".csv"):
        file_path = os.path.join("日志数据/混淆矩阵", file_name)
        acc_total, acc_per_class, recall_per_class, f1_per_class = get_confusion_all(file_path)
        # 模型名为当前csv文件名去除confusion_matrix
        model_name = file_name.replace("_confusion_matrix.csv", "")
        print(
            f"模型：{model_name}：acc_total:{acc_total}  F1:{f1_per_class}  ACC:{acc_per_class}  Recall:{recall_per_class}")
        # 写入文件txt
        with open("日志数据/测试/评价结果.txt", "a") as f:
            f.write(
                f"模型：{model_name}：acc_total:{acc_total}  F1:{f1_per_class}  ACC:{acc_per_class}  Recall:{recall_per_class}\n")
